#!/usr/bin/env python
from __future__ import print_function
import os
import numpy as np
from maputility2 import hpxtools, maptools
from XmlTools import SMKernel
from glat2 import mdltk
from glat2.likeold import NewLikeObj

normpars = set(['Prefactor', 'norm', 'Integral', 'Value', 'Normalization'])

def parseXml_allfree(xmlfile):
    sm = SMKernel.SourceModel(xmlfile, raiseErr=True)
    print('[parseXml_allfree] input xmlfile: %s' % xmlfile)

    freesrcs = {}
    for srcname, src in sm.srcList.items():
        srcpars, miscpars, freepars = {}, {}, {}

        spectype = src.spectrum.type
        if spectype == 'FileFunction':
            srcpars['file'] = src.spectrum.file

        miscpars['apply_edisp'] = True
        if 'apply_edisp' in src.spectrum.__dict__:
            apply_edisp_lower = src.spectrum.apply_edisp.lower()
            if (apply_edisp_lower == 'true') or (apply_edisp_lower == '1') or \
                    (apply_edisp_lower == 't'):
                miscpars['apply_edisp'] = True
            else:
                miscpars['apply_edisp'] = False

        if 'scaling_file' in src.spectrum.__dict__:
            miscpars['aeff_scalefile'] = src.spectrum.scaling_file
        else:
            miscpars['aeff_scalefile'] = None

        for pname, parobj in src.spectrum.parameters.items():
            srcpars[pname] = parobj.value * parobj.scale
            if pname in normpars:
                freepars[pname] = (parobj.value, parobj.min, parobj.max, parobj.scale, None, None)

        freesrcs[srcname] = (spectype, srcpars, freepars, miscpars)

    return freesrcs, {}


def run_dpmodel(srcmap, mdlfile, outmcube, healpix=False,
        minbin=None, maxbin=None, ebdmin=None, ebdmax=None, save_comps=True, drmfile=None):
    l = NewLikeObj(None, None) # Manual Mode
    frees, fixs = parseXml_allfree(mdlfile)
    l.like = l._initChiSqObj(frees, fixs, srcmap, isHealpix=healpix, drmfile=drmfile)

    #l._init_Minuit((minbin, maxbin))
    l._init_Minuit()
    l.like()

    myminbin, mymaxbin = minbin, maxbin
    if myminbin is None:
        if ebdmin is not None:
            is_close = np.isclose(l.like.energies, ebdmin)
            assert is_close.sum() == 1 # only one energy bin is close to the ebdmin
            myminbin = int(np.nonzero(is_close)[0][0])
        else:
            myminbin = 0

    if mymaxbin is None:
        if ebdmax is not None:
            is_close = np.isclose(l.like.energies, ebdmax)
            assert is_close.sum() == 1 # only one energy bin is close to the ebdmax
            mymaxbin = int(np.nonzero(is_close)[0][0]) - 1
        else:
            mymaxbin = l.like.energies.size-1

    l.like.toModel(outmcube, minbin=myminbin, maxbin=mymaxbin, save_comps=save_comps)
    print('.'*30, 'Done!', '.'*30)

    return l, myminbin, mymaxbin


def merge_mcube(mcubefiles, outfile, healpix=False, extidx=None):
    """
    mcubefiles is a list and assumed to arranged with increasing energies
    """
    # collect data
    ebins, mapdata = [], []
    for mcubefile in mcubefiles:
        if healpix:
            m = hpxtools.hpxOpen(mcubefile, extidx=extidx)
        else:
            m = maptools.mapOpen(mcubefile, extidx=extidx)

        ebobj = mdltk.EnergyBins(mcubefile)
        if ebins:
            assert np.isclose(ebobj.energies_MeV[0], ebins[-1]) # assume Ebins for two files are connected
            idxlo = len(ebins) - 1
            ebins.extend(ebobj.energies_MeV[1:])
        else:
            idxlo = 0
            ebins.extend(ebobj.energies_MeV)
        idxhi = len(ebins)-1

        mapdata.append((m.mapdata, idxlo, idxhi))

    # merge
    nebins = len(ebins)
    if nebins>0: # make sure mcubefile is not empty
        if healpix:
            mnew = hpxtools.hpxLike(mcubefiles[0], nlayers=nebins-1)
        else:
            mnew = maptools.mapLike(mcubefiles[0], nlayers=nebins-1)

        for mdata, idxlo, idxhi in mapdata:
            print(idxlo, idxhi)
            mnew.mapdata[idxlo:idxhi] = mdata

    # save
    mnew.saveMap(outfile)

    ebnew = mdltk.EnergyBins(ebins)
    ebnew.append_hdu_tofile(outfile, ebtype=1)


def get_mcube_universal(srcmap, mdlfile_first, outmcube, healpix=False, drmfile=None):
    path, srcmdl = os.path.split(mdlfile_first)
    if not path: path = './'
    valid_suffix = '.'+srcmdl.split('.')[-1]

    files = []
    for f in os.listdir(path):
        if f.endswith(valid_suffix):
            files.append(f)

    is_sed = False
    mcube, ebins = None, None
    while srcmdl is not None:
        fn0 = srcmdl.split('.')[0]
        if 'wid' in fn0:
            fn0, suff = fn0.split('_wid')
            widbin = int(suff)
        else:
            widbin = 1

        nextsrcmdl = None
        srcmdl_full = os.path.join(path, srcmdl)
        if 'band' in fn0:
            is_sed = True
            pref, suff = fn0.split('band')
            idxmin = int(suff)
            idxmax = idxmin + widbin

            pref_next = pref+'band%02d'%idxmax
            for f in files:
                if f.startswith(pref_next):
                    nextsrcmdl = f

            l, ilo, ihi = run_dpmodel(srcmap, srcmdl_full, outmcube, healpix, minbin=idxmin, maxbin=idxmax-1,
                    save_comps=False, drmfile=drmfile)
        else:
            # for broad band
            run_dpmodel(srcmap, srcmdl_full, outmcube, healpix, save_comps=False, drmfile=drmfile)
            break

        if healpix:
            mcube0 = hpxtools.hpxOpen(outmcube)
        else:
            mcube0 = maptools.mapOpen(outmcube)

        if mcube is None:
            datfile, extidx= l.like.data.mapData.filename, l.like.data.mapData.extidx
            ebins = mdltk.EnergyBins(datfile)
            # nan to ensure not to get invalid data outside the energy range
            if healpix:
                mcube = hpxtools.hpxLike(datfile, extidx=extidx, fill_value=np.nan)
            else:
                mcube =maptools.mapLike(datfile, extidx=extidx, fill_value=np.nan)

        mcube.mapdata[ilo:ihi+1, :, :] = mcube0.mapdata[:,:,:]
        srcmdl = nextsrcmdl

    if is_sed:
        mcube.saveMap(outmcube)
        ebins.append_hdu_tofile(outmcube)


if __name__ == '__main__':
    import argparse
    parser = argparse.ArgumentParser(prefix_chars='-+',
            description='a quick tool to get the model cube given the srcmaps and xmlfile')
    parser.add_argument("srcmap", type=str, help='srcmap')
    parser.add_argument("mdlfile", type=str,
            help='input xml/yaml model file, no bash variable allowed')
    parser.add_argument("-drmfile", type=str, default=None,
            help='The name of the drm file')
    parser.add_argument("-outmcube", type=str, default='mcube_fitted.fits',
            help='The name of output mcube file (default: mcube_fitted.fits)')
    parser.add_argument("-minbin", type=int, default=None,
            help='The bin range [minbin:maxbin] will be included (default: 0)')
    parser.add_argument("-maxbin", type=int, default=None,
            help='The bin range [minbin:maxbin] will be included (default: MAX_NLAYER)')
    parser.add_argument("+healpix", action='store_true', default=False,
            help='whether the srcmaps are healpix')
    parser.add_argument("-comps", action='store_false', default=True,
            help='whether to save the gamma-ray components in the fits (default: True)')
    parser.add_argument("+sed", action='store_true', default=False,
            help='whether to deal with sed data (comps is not valid)')
    args = parser.parse_args()

    if args.sed:
        get_mcube_universal(srcmap=args.srcmap, mdlfile_first=args.mdlfile, outmcube=args.outmcube,
                healpix=args.healpix, drmfile=args.drmfile)
    else:
        run_dpmodel(srcmap=args.srcmap, mdlfile=args.mdlfile, outmcube=args.outmcube, healpix=args.healpix,
            minbin=args.minbin, maxbin=args.maxbin, save_comps=args.comps, drmfile=args.drmfile)
